use bytes::Bytes;
use futures::future::poll_fn;
use futures::sink::SinkExt;
use futures::stream::TryStreamExt;
use futures::{Sink, Stream};
use once_cell::sync::Lazy;
use prometheus::{register_histogram, Histogram};
use std::net::SocketAddr;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
use tokio::net::{TcpStream, ToSocketAddrs};
use tokio::time::timeout;
use tokio_util::codec::Framed;

use crate::codec::FloFrameCodec;
use crate::error::*;
use crate::packet::{FloPacket, Frame};
use tokio::io::AsyncWriteExt;

const DEFAULT_TIMEOUT: Duration = Duration::from_secs(3);

pub static PACKET_PROCESSING_DURATION: Lazy<Histogram> = Lazy::new(|| {
  register_histogram!(
    "flonet_w3gs_packet_processing_duration_seconds",
    "Duration of W3GS packet processing in seconds",
    vec![
      0.0001, 0.00025, 0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5,
      5.0, 10.0
    ]
  )
  .unwrap()
});

#[derive(Debug)]
pub struct FloStream {
  pub timeout: Duration,
  pub(crate) transport: Framed<TcpStream, FloFrameCodec>,
}

impl FloStream {
  pub async fn connect_no_delay<A: ToSocketAddrs>(addr: A) -> Result<Self> {
    let socket = timeout(DEFAULT_TIMEOUT, TcpStream::connect(addr))
      .await
      .map_err(|_elapsed| Error::StreamTimeout)??;

    socket.set_nodelay(true).ok();

    //TODO: not supported by current tokio
    //socket.set_keepalive(None).ok();

    let transport = Framed::new(socket, FloFrameCodec::new());
    Ok(FloStream {
      transport,
      timeout: DEFAULT_TIMEOUT,
    })
  }

  pub async fn connect<A: ToSocketAddrs>(addr: A) -> Result<Self> {
    let socket = timeout(DEFAULT_TIMEOUT, TcpStream::connect(addr))
      .await
      .map_err(|_elapsed| Error::StreamTimeout)??;

    // not supported by tokio atm
    //socket.set_keepalive(Some(Duration::from_secs(30)))?;

    let transport = Framed::new(socket, FloFrameCodec::new());
    Ok(FloStream {
      transport,
      timeout: DEFAULT_TIMEOUT,
    })
  }

  pub fn new(socket: TcpStream) -> Self {
    FloStream {
      transport: Framed::new(socket, FloFrameCodec::new()),
      timeout: DEFAULT_TIMEOUT,
    }
  }

  pub fn set_timeout(&mut self, duration: Duration) -> &mut Self {
    self.timeout = duration;
    self
  }

  #[inline]
  pub fn local_addr(&self) -> Result<SocketAddr> {
    self.transport.get_ref().local_addr().map_err(Into::into)
  }

  #[inline]
  pub fn peer_addr(&self) -> Result<SocketAddr> {
    self.transport.get_ref().peer_addr().map_err(Into::into)
  }

  pub async fn send_frame_timeout(&mut self, frame: Frame) -> Result<()> {
    timeout(self.timeout, self.transport.send(frame))
      .await
      .map_err(|_elapsed| Error::StreamTimeout)??;
    Ok(())
  }

  #[inline]
  pub async fn send_frame(&mut self, frame: Frame) -> Result<()> {
    let original_action_received_at = frame.correlation_timestamp;
    self.transport.send(frame).await?;
    if let Some(original_action_received_at) = original_action_received_at {
      PACKET_PROCESSING_DURATION.observe(
        std::time::Instant::now()
          .duration_since(original_action_received_at)
          .as_secs_f64(),
      );
    }
    Ok(())
  }

  #[inline]
  pub async fn send_frames<I>(&mut self, iter: I) -> Result<()>
  where
    I: IntoIterator<Item = Frame>,
  {
    let mut stream = tokio_stream::iter(iter.into_iter().map(Ok));
    timeout(self.timeout, self.transport.send_all(&mut stream))
      .await
      .map_err(|_elapsed| Error::StreamTimeout)??;
    Ok(())
  }

  #[inline]
  pub async fn send<T>(&mut self, packet: T) -> Result<()>
  where
    T: FloPacket,
  {
    self.send_frame_timeout(packet.encode_as_frame()?).await?;
    Ok(())
  }

  #[inline]
  pub async fn recv<T>(&mut self) -> Result<T>
  where
    T: FloPacket + Default,
  {
    let frame = self.recv_frame().await?;
    Ok(frame.decode()?)
  }

  #[inline]
  pub async fn recv_timeout<T>(&mut self, duration: Duration) -> Result<T>
  where
    T: FloPacket + Default,
  {
    let frame = timeout(duration, self.recv_frame())
      .await
      .map_err(|_elapsed| Error::StreamTimeout)??;
    Ok(frame.decode()?)
  }

  #[inline]
  pub async fn recv_frame(&mut self) -> Result<Frame> {
    let frame = self
      .transport
      .try_next()
      .await?
      .ok_or_else(|| Error::StreamClosed)?;
    Ok(frame)
  }

  #[inline]
  pub async fn recv_frame_timeout(&mut self) -> Result<Frame> {
    let frame = timeout(self.timeout, self.transport.try_next())
      .await
      .map_err(|_elapsed| Error::StreamTimeout)??
      .ok_or_else(|| Error::StreamClosed)?;
    Ok(frame)
  }

  pub async fn flush(&mut self) -> Result<()> {
    poll_fn(|ctx| Pin::new(&mut self.transport).poll_flush(ctx)).await?;
    self.transport.get_mut().flush().await?;
    Ok(())
  }

  pub async fn shutdown(&mut self) -> Result<()> {
    poll_fn(|ctx| Pin::new(&mut self.transport).poll_close(ctx)).await?;
    self.transport.get_mut().shutdown().await?;
    Ok(())
  }

  pub async fn downgrade_to_binary_stream(self) -> Result<(Bytes, TcpStream)> {
    let parts = self.transport.into_parts();
    let mut stream = parts.io;
    if !parts.write_buf.is_empty() {
      stream.write_all(parts.write_buf.as_ref()).await?;
    }
    Ok((parts.read_buf.freeze(), stream))
  }
}

impl Stream for FloStream {
  type Item = Result<Frame>;

  fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
    Pin::new(&mut self.transport).poll_next(cx)
  }
}

#[test]
fn test_lookup() {
  use std::net::ToSocketAddrs;
  let mut addrs_iter = "wc3.tools:443".to_socket_addrs().unwrap();
  dbg!(addrs_iter.next());
}
